import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Union
import logging

logger = logging.getLogger(__name__)

class Portfolio:
    """
    管理账户的资产状况。
    """
    def __init__(self, initial_cash: float = 1000000.0):
        self.initial_cash = initial_cash
        self.cash = initial_cash
        self.positions = {}  # {asset: {'shares': number, 'cost_price': price, 'last_price': price}}
        self.equity_curve = []
        self.orders = {}
        
        # 历史记录
        self.history = {
            'cash': [initial_cash],
            'total_value': [initial_cash],
            'positions_value': [0.0],
            'unrealized_pnl': [0.0],
            'realized_pnl': 0.0
        }
        
        # 风险管理参数
        self.max_position_size = 0.1  # 单个持仓最大比例
        self.max_total_leverage = 1.0  # 最大杠杆率
        
        logger.info(f"Portfolio initialized with cash: {initial_cash:,.2f}")
    
    def update_total_value(self, current_prices: Dict[str, float]) -> float:
        """
        根据当前价格计算总资产价值
        """
        try:
            positions_value = 0.0
            unrealized_pnl = 0.0
            
            for asset, position in self.positions.items():
                if asset in current_prices:
                    shares = position['shares']
                    current_price = current_prices[asset]
                    cost_price = position['cost_price']
                    
                    # 更新最新价格
                    position['last_price'] = current_price
                    
                    # 计算市值
                    market_value = shares * current_price
                    positions_value += market_value
                    
                    # 计算未实现损益
                    unrealized_pnl += shares * (current_price - cost_price)
                else:
                    logger.warning(f"No current price available for {asset}")
                    # 使用上次价格
                    last_price = position.get('last_price', position['cost_price'])
                    positions_value += position['shares'] * last_price
            
            total_value = self.cash + positions_value
            
            # 更新历史记录
            self.history['cash'].append(self.cash)
            self.history['total_value'].append(total_value)
            self.history['positions_value'].append(positions_value)
            self.history['unrealized_pnl'].append(unrealized_pnl)
            
            return total_value
            
        except Exception as e:
            logger.error(f"Error updating portfolio value: {str(e)}")
            return self.cash  # 返回现金作为备用
    
    def get_position_info(self, asset: str) -> Dict:
        """获取持仓信息"""
        if asset not in self.positions:
            return {'shares': 0, 'cost_price': 0, 'last_price': 0, 'market_value': 0, 'unrealized_pnl': 0}
        
        position = self.positions[asset]
        shares = position['shares']
        cost_price = position['cost_price']
        last_price = position.get('last_price', cost_price)
        
        return {
            'shares': shares,
            'cost_price': cost_price,
            'last_price': last_price,
            'market_value': shares * last_price,
            'unrealized_pnl': shares * (last_price - cost_price),
            'unrealized_pnl_pct': (last_price - cost_price) / cost_price if cost_price > 0 else 0
        }
    
    def get_total_positions_value(self) -> float:
        """获取所有持仓的总市值"""
        total_value = 0.0
        for asset, position in self.positions.items():
            last_price = position.get('last_price', position['cost_price'])
            total_value += position['shares'] * last_price
        return total_value
    
    def get_available_cash(self) -> float:
        """获取可用现金"""
        return max(0, self.cash)
    
    def check_buying_power(self, amount: float) -> bool:
        """检查购买力"""
        return self.get_available_cash() >= amount
    
    def calculate_position_limit(self, asset: str, current_price: float) -> int:
        """计算最大可购买数量"""
        try:
            total_value = self.cash + self.get_total_positions_value()
            max_position_value = total_value * self.max_position_size
            max_shares = int(max_position_value / current_price)
            
            # 考虑现金限制
            cash_limit_shares = int(self.cash / current_price)
            
            return min(max_shares, cash_limit_shares)
            
        except Exception as e:
            logger.error(f"Error calculating position limit for {asset}: {str(e)}")
            return 0
    
    def get_portfolio_summary(self) -> Dict:
        """获取投资组合概要"""
        total_positions_value = self.get_total_positions_value()
        total_value = self.cash + total_positions_value
        
        return {
            'cash': self.cash,
            'positions_value': total_positions_value,
            'total_value': total_value,
            'cash_ratio': self.cash / total_value if total_value > 0 else 1.0,
            'positions_ratio': total_positions_value / total_value if total_value > 0 else 0.0,
            'total_return': (total_value - self.initial_cash) / self.initial_cash,
            'total_return_pct': ((total_value - self.initial_cash) / self.initial_cash) * 100,
            'num_positions': len(self.positions),
            'realized_pnl': self.history['realized_pnl']
        }
    
    def get_positions_summary(self) -> List[Dict]:
        """获取所有持仓的详细信息"""
        summary = []
        for asset, position in self.positions.items():
            info = self.get_position_info(asset)
            info['asset'] = asset
            summary.append(info)
        
        # 按市值排序
        summary.sort(key=lambda x: x['market_value'], reverse=True)
        return summary
    
    def validate_order(self, asset: str, shares: int, price: float) -> bool:
        """验证订单是否有效"""
        if shares == 0:
            return False
        
        if shares > 0:  # 买入
            required_cash = shares * price
            if not self.check_buying_power(required_cash):
                logger.warning(f"Insufficient cash for buying {shares} shares of {asset}")
                return False
        else:  # 卖出
            current_shares = self.positions.get(asset, {}).get('shares', 0)
            if abs(shares) > current_shares:
                logger.warning(f"Insufficient shares for selling {abs(shares)} shares of {asset}")
                return False
        
        return True
    
    def add_realized_pnl(self, amount: float):
        """添加已实现损益"""
        self.history['realized_pnl'] += amount
